//
// Created by yk120 on 2024/3/4.
//
#include <nanofaiss/IndexFlat.h>
#include <nanofaiss/IndexLightWeight.h>
#include <nanofaiss/utils/utils.h>
#include <random>
#include <unordered_set>

// 64-bit int
using idx_t = faiss::idx_t;
using faiss::getmillisecs;


// belong to 1 2 3 4
// sample rate 0.01  0.05 0.1 0.5

// x is the query result, y is the ground truth
float calculate_recall_k(idx_t *x, idx_t *y, int k) {
    float number = 0.f;
    std::unordered_set<idx_t> set;
    for (size_t i = 0; i != k; i++) set.insert(x[i]);
    for (size_t i = 0; i != k; i++)
        if (set.count(y[i])) number += 1;
    return number / k;
}

int main() {
    int d = 128;      // dimension
    int nb = 10000; // database size
    int nq = 1000;  // nb of queries

    std::mt19937 rng;
    std::uniform_real_distribution<> distrib;

    float* xb = new float[d * nb];
    float* xq = new float[d * nq];

    for (int i = 0; i < nb; i++) {
        for (int j = 0; j < d; j++)
            xb[d * i + j] = distrib(rng);
        xb[d * i] += i / 1000.;
    }

    for (int i = 0; i < nq; i++) {
        for (int j = 0; j < d; j++)
            xq[d * i + j] = distrib(rng);
        xq[d * i] += i / 1000.;
    }

    int k = 10; // TOP-K

    idx_t* I1 = new idx_t[k * nq];
    float* D1 = new float[k * nq];

    idx_t* I2 = new idx_t[k * nq];
    float* D2 = new float[k * nq];

    printf("###############FlatIndex##############\n");
    faiss::IndexFlatL2 index(d); // call constructor
    printf("is_trained = %s\n", index.is_trained ? "true" : "false");
    index.add(nb, xb); // add vectors to the index
    printf("ntotal = %zd\n", index.ntotal);
    double t1 = getmillisecs();
    index.search(nq, xq, k, D2, I2);
    double t2 = getmillisecs();
    printf("Flat Search Time: %lf\n", t2  - t1);

    printf("###############IndexLightWeight##############\n");
    faiss::IndexLightWeight indexL(d, 0.1);
    printf("is_trained = %s\n", indexL.is_trained ? "true" : "false");
    indexL.add(nb, xb); // add vectors to the index
    indexL.train(2);
    printf("ntotal = %zd\n", indexL.ntotal);

    delete[] xb;



    {

        {
            for (size_t probe = 1; probe < 200; probe += 10) {
                double t1 = getmillisecs();
                indexL.search(xq, nq, k, D1, I1, probe);
                double t2 = getmillisecs();
                double td = t2 - t1;
                printf("LWI Search Time: %lf, probe = %lu\n", t2 - t1, probe);


                // compute recall
                float recall_k = 0;
                for (size_t i = 0; i < nq; i++) {
                    idx_t *p = I1 + k * i;
                    idx_t *g = I2 + k * i;
                    float current_recall = calculate_recall_k(p, g, k);
                    recall_k += current_recall;
                }
                printf("(Search Time, Average recall) = (%lf, %f)\n", td, recall_k / nq);
            }
        }

    }





    delete[] I1;
    delete[] D1;

    delete[] I2;
    delete[] D2;


    delete[] xq;

    return 0;
}
